# 这是ocr的第二个版本 实现对题型的自动判别以及相应的后续处理（阅读，选择，翻译，作文）
# 通用 返回题型，ocr结果，在es中的id
# 阅读：领域 关键词
# 选择：每个选择题的考点
# 作文: 关键词 领域
# 翻译: 关键词 领域
# 功能：对（选择题，阅读题，翻译题，作文题）实现ocr,并将其返回为正确的格式
import requests
import base64
import re
import os
import json
from sparkai.llm.llm import ChatSparkLLM, ChunkPrintHandler
from sparkai.core.messages import ChatMessage
import threading
import time
from prompt.choice_prompt import choice_prompt
from prompt.reading_prompt import reading_prompt
from prompt.identify_prompt import identify_prompt
import es_option.es_api as es_option
from es_option.es_api import get_gpt_embedding
import gpt.gpt_key as gpt_key
import sys

# start_time = time.time()
access_token = "24.9b0dd814122ce08ebafd1a4ada96d3fb.2592000.1716464568.282335-56894520"
api_key = gpt_key.gpt_key
os.environ["http_proxy"] = "http://127.0.0.1:10808"
os.environ["https_proxy"] = "http://127.0.0.1:10808"

SPARKAI_URL = 'wss://spark-api.xf-yun.com/v3.1/chat'
SPARKAI_APP_ID = '1d38b17a'
SPARKAI_API_SECRET = 'MTJjOGQ5YzE1ZWZjYjJkOWUxZTM4Mzcx'
SPARKAI_API_KEY = 'b33d686049c006fd12a7e9b345ec0409'
SPARKAI_DOMAIN = 'generalv3'

spark = ChatSparkLLM(
    spark_api_url=SPARKAI_URL,
    spark_app_id=SPARKAI_APP_ID,
    spark_api_key=SPARKAI_API_KEY,
    spark_api_secret=SPARKAI_API_SECRET,
    spark_llm_domain=SPARKAI_DOMAIN,
    streaming=False,
)

# 用于线程返回结果的变量
res_reading = None  # 都修改成了json
res_type = None  # 都修改成了json
res_choice = None  # 都修改成了json

# 实现ocr输出 并且矫正了格式
def ocr_first(baidu_api, src):
    """
    实现对不同题型拍照
    :param baidu_api: 百度ocr_key
    :param src: 图片的路径
    :return: 对应的json格式，不同题，不同json
    """
    access_token = baidu_api
    request_url = "https://aip.baidubce.com/rest/2.0/ocr/v1/accurate"
    f = open(src, "rb")
    img = base64.b64encode(f.read())
    params = {"image": img, "language_type": "CHN_ENG", "result_type": "big"}
    request_url = request_url + "?access_token=" + access_token
    headers = {"content-type": "application/x-www-form-urlencoded"}
    response = requests.post(request_url, data=params, headers=headers)
    print(response.json())
    return response

def ocr_srcond(response, type):
    """
    此时已经获得了ocr的返回的初步内容，然后对其进行格式微调
    :param response: ocr返回的初步结果
    :param type: 类型 reading choice translation writing
    :return:
    """
    if type == "reading":
        max_width = 0
        for words in response.json()["words_result"]:
            max_width = max(max_width, words["location"]["width"])
        article = " "
        for words in response.json()["words_result"]:
            if words["location"]["width"] >= max_width - 50:
                article = article + words["words"] + " "
            else:
                article = article + words["words"] + "\n" + "  "
        return article
    if type == "choice":
        max_width = 0
        for words in response.json()["words_result"]:
            max_width = max(max_width, words["location"]["width"])
        choice_partten = "^[abcdABCD]\..*"  # 以A. B.等等开头的
        article = " "
        now_words = response.json()["words_result"][0]
        same_row = 0
        for i, words in enumerate(response.json()["words_result"]):
            if i > 0:
                if (
                    now_words["location"]["top"] + 10
                    > words["location"]["top"]
                    > now_words["location"]["top"] - 10
                ):
                    same_row = same_row + 1
                    if re.match(choice_partten, words["words"]):
                        article = article + now_words["words"] + " "
                        same_row = 0
                    else:
                        article = article + now_words["words"] + "( )"
                else:
                    if same_row == 0 and not re.match(
                        choice_partten, now_words["words"]
                    ):
                        article = article + now_words["words"] + "( )"
                    else:
                        article = article + now_words["words"] + " "
                now_words = words
        article = article + response.json()["words_result"][-1]["words"]
        return article
    if type == "translation" or type == "writing":
        max_width = 0
        for words in response.json()["words_result"]:
            max_width = max(max_width, words["location"]["width"])
        article = " "
        for words in response.json()["words_result"]:
            if words["location"]["width"] >= max_width - 50:
                article = article + words["words"] + " "
            else:
                article = article + words["words"] + "\n" + "  "
        return article

def call_sparkai(prompt):
    """
    封装调用星火大模型的接口
    :param prompt: 输入给星火大模型的提示词语
    :return: str 星火大模型的返回结果
    """
    messages = [ChatMessage(role="user", content=prompt)]
    handler = ChunkPrintHandler()
    response = spark.generate([messages], callbacks=[handler])
    return response

# 对于判断题型的函数
def identify_type(ocr_res, api_key):
    """
    实现对题型的判定
    :param ocr_res: 初步的ocr结果,百度云的结果
    :return:Example JSON format:
    {
    "Type": "choice"
    }
    {
    "Type": "reading",
    "Filed": "The field to which the reading question belongs",
    "Keyword": ["The keywords for this reading question"]
    }
    {
    "Type": "translation",
    "Filed": "The field to which the translation question belongs",
    "Keyword": ["The keywords for this reading question"]
    }
    {
    "Type": "writing",
    "Filed": "The field to which the essay question belongs",
    "Keyword": ["The keywords for this essay question"],
    "Genre": "The genre of this essay"
    }
    """
    global res_type  # 最后需要这个结果
    ocr_res = ocr_srcond(ocr_res, "reading")
    identify_ocr_prompt = identify_prompt.format(ocr_res)
    res = call_sparkai(identify_ocr_prompt)
    res_type = json.loads(res)

# 对于阅读部分的最终结果
def reading_ocr_res(ocr_res, spark_key):
    """
    对ocr的结果，进行修改，返回给前端,writing和translation返回正确格式的字符串,选择和阅读，则是字符串，
    :param baidu_api: 百度云key
    :param spark_key: gptkey
    :param src: 图片路径
    :param type: 题目类型
    :return:
    """
    global res_reading
    ocr_res = ocr_srcond(ocr_res, "reading")
    reading_ocr_prompt = reading_prompt.format(ocr_res)
    result = call_sparkai(reading_ocr_prompt)
    res_reading = json.loads(result)
    return res_reading

# 对于选择题的部分的最终结果
def choice_ocr_res(ocr_res, spark_key):
    """
    对于选择题部分的ocr最后结果
    :param ocr_res: 初次访问百度云的结果
    :param spark_key:
    :return: 对应的json格式
     "questions": [
    {
      "question": "What is the capital of France?",
      "options": ["Paris", "Berlin", "London", "Rome"],
     "knowledgePoints": "The knowledge point to which this multiple-choice question belongs"
    }]
    """
    global res_choice
    ocr_res = ocr_srcond(ocr_res, "choice")
    prompt = choice_prompt.format(ocr_res)
    res = call_sparkai(prompt)
    res_choice = json.loads(res)

# 最后的ocr过程:
def final_ocr(baidu_api, spark_key, src):
    # 创建线程
    ocr_res = ocr_first(baidu_api=baidu_api, src=src)  # 返回的就是百度云返回的结果
    reading_thread = threading.Thread(target=reading_ocr_res, args=[ocr_res, spark_key])
    identify_thread = threading.Thread(target=identify_type, args=[ocr_res, spark_key])
    choice_thread = threading.Thread(target=choice_ocr_res, args=[ocr_res, spark_key])
    # 启动线程
    reading_thread.start()
    identify_thread.start()
    choice_thread.start()
    res_translation = ocr_srcond(ocr_res, "translation")
    res_writing = ocr_srcond(ocr_res, "writing")
    # 等待线程结束
    identify_thread.join()
    if res_type["type"] == "reading":
        reading_thread.join()
        # 处理阅读题的逻辑
        # ...
    elif res_type["type"] == "choice":
        choice_thread.join()
        # 处理选择题的逻辑
        # ...
    elif res_type["type"] == "translation":
        # 处理翻译题的逻辑
        # ...
    elif res_type["type"] == "writing":
        # 处理作文题的逻辑
        # ...

if __name__ == "__main__":
    src = "yuedu.jpg"
    src = os.path.abspath(".") + "/api/target/" + src
    final_ocr(baidu_api=access_token, spark_key=api_key, src=src)
